"""ONNX Runtime 适配器实现"""

from __future__ import annotations

import os
import sys
from pathlib import Path
from typing import Any, Dict, List

# 添加 NVIDIA CUDA 库路径到 PATH（用于 Windows pip 安装的 CUDA）
def _setup_cuda_paths():
    """将 pip 安装的 NVIDIA CUDA 库添加到 PATH"""
    if sys.platform == "win32":
        site_packages = Path(sys.prefix) / "Lib" / "site-packages"
        nvidia_paths = [
            site_packages / "nvidia" / "cudnn" / "bin",
            site_packages / "nvidia" / "cublas" / "bin",
            site_packages / "nvidia" / "cuda_runtime" / "bin",
        ]
        for path in nvidia_paths:
            if path.exists():
                path_str = str(path)
                if path_str not in os.environ.get("PATH", ""):
                    os.environ["PATH"] = path_str + os.pathsep + os.environ.get("PATH", "")

_setup_cuda_paths()

import numpy as np

from .base import AdapterBackend
from ...exceptions import ModelLoadError, InferenceError, ComplianceError


class OnnxRuntimeAdapter(AdapterBackend):
    """ONNX Runtime 后端适配器，符合 adapter-spec.md 规范"""

    def __init__(self):
        self._session = None

    @staticmethod
    def is_available() -> bool:
        """
        检查 ONNX Runtime 是否可用
        
        Returns:
            bool: 如果 onnxruntime 已安装则返回 True
        """
        try:
            import onnxruntime
            return True
        except ImportError:
            return False

    def prepare(self, model_spec: Dict[str, Any]) -> Any:
        """
        加载 ONNX 模型并创建推理会话
        
        Args:
            model_spec: 模型规格字典，包含 source_onnx, sha256 等字段
            
        Returns:
            onnxruntime.InferenceSession 推理会话对象
            
        Raises:
            ModelLoadError: 模型加载失败
            ComplianceError: 模型哈希校验失败
        """
        try:
            import onnxruntime as ort
            
            model_path = model_spec.get("source_onnx")
            if not model_path:
                raise ModelLoadError("Missing 'source_onnx' in model_spec")
            
            # 检查文件是否存在
            if not Path(model_path).exists():
                raise ModelLoadError(f"Model file not found: {model_path}")
            
            # 可选：校验模型哈希（如果提供）
            self.validate_model_compliance(model_spec)
            
            # 配置会话选项
            sess_options = ort.SessionOptions()
            sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
            
            # 选择执行提供者（仅使用 CUDA，不使用 CPU 回退）
            # 注意：如果 CUDA 不可用或某些算子不支持 CUDA，推理会失败
            providers = ["CUDAExecutionProvider"]
            
            # 创建推理会话
            session = ort.InferenceSession(
                model_path,
                sess_options=sess_options,
                providers=providers
            )
            
            # 检查实际使用的 providers
            actual_providers = session.get_providers()
            primary_provider = actual_providers[0] if actual_providers else "Unknown"
            print(f"[INFO] 主执行器: {primary_provider}")
            
            self._session = session
            return session
            
        except Exception as e:
            if isinstance(e, (ModelLoadError, ComplianceError)):
                raise
            raise ModelLoadError(f"Failed to load ONNX model: {e}") from e

    def infer(self, handle: Any, inputs: Dict[str, np.ndarray]) -> List[np.ndarray]:
        """
        执行单次推理
        
        Args:
            handle: prepare() 返回的推理会话
            inputs: 输入张量字典 {input_name: np.ndarray}
            
        Returns:
            输出张量列表 List[np.ndarray]
            
        Raises:
            InferenceError: 推理执行失败
        """
        try:
            outputs = handle.run(None, inputs)
            return outputs
        except Exception as e:
            raise InferenceError(f"Inference failed: {e}") from e

    def teardown(self, handle: Any) -> None:
        """
        释放推理句柄和相关资源
        
        Args:
            handle: prepare() 返回的推理会话
        """
        # ONNX Runtime 会话由 Python GC 自动管理，无需显式清理
        # 注意：保留 self._session 引用，以便后续 get_metadata() 可以获取实际使用的 providers
        pass

    def get_metadata(self) -> Dict[str, Any]:
        """
        返回适配器和后端的元信息
        
        Returns:
            包含适配器名称、版本、后端信息等的字典
        """
        import onnxruntime as ort
        
        # 获取当前会话使用的 providers，如果会话不存在则返回可用的 providers
        if self._session is not None:
            providers = self._session.get_providers()
        else:
            providers = ort.get_available_providers()
        
        device = "CUDA" if "CUDAExecutionProvider" in providers else "CPU"
        
        # 尝试获取 CUDA 设备信息
        cuda_version = None
        if "CUDAExecutionProvider" in providers:
            try:
                from nvidia_ml_py import pynvml
                pynvml.nvmlInit()
                handle = pynvml.nvmlDeviceGetHandleByIndex(0)
                cuda_version = pynvml.nvmlSystemGetCudaDriverVersion()
                pynvml.nvmlShutdown()
            except Exception:
                pass
        
        return {
            "adapter_backend": "onnxruntime",
            "adapter_version": "1.0.0",
            "backend": "onnxruntime",
            "backend_version": ort.__version__,
            "device": device,
            "providers": providers,
            "cuda_version": cuda_version,
            "conversion_tool": None,  # 直接使用 ONNX，无需转换
            "optimization_flags": ["graph_opt_all"]
        }

    @staticmethod
    def _print_node_placement(session):
        """打印模型算子在各个 EP 上的分配情况"""
        try:
            # 获取会话信息（需要启用 profiling 或使用内部 API）
            # 这里我们使用一个简单的方法：统计 providers
            providers = session.get_providers()
            print(f"[INFO] 算子分配策略: 优先使用 {providers[0]}")
            
            # 注意：ONNX Runtime 不提供直接的 API 来查看每个算子的分配
            # 如果需要详细信息，需要启用 verbose logging
            print("[INFO] 如需查看详细的算子分配，请设置环境变量：")
            print("       export ORT_LOG_SEVERITY_LEVEL=1")
            print("       export ORT_DISABLE_ALL_OPTIMIZATION=1")
        except Exception as e:
            print(f"[DEBUG] 无法获取算子分配详情: {e}")
    

